﻿using System;
using System.Collections.Generic;
using System.Linq;
using Algoritmia.StateSpaces;
using Algoritmia.Utils;
using Lists = Algoritmia.DataStructures.Lists;

namespace Algoritmia.Problems.Puzzles.NQueens
{
    public class NQueensStateSpace1 : IForwardStateSpace<Lists.IList<int>,int>
    {
        private int _n;

        public NQueensStateSpace1(int N)
        {
            _n = N;
        }

        public Lists.IList<int> Decide(Lists.IList<int> state, int decision)
        {
            var list = AlgoritmiaFactory.Lists<int>.BuildDefault(state);
            list.Add(decision);
            return list;
        }

        public IEnumerable<int> Decisions(Lists.IList<int> state)
        {
            for(int row = 0; row < _n; row++)
            {
                if (!state.Contains(row) && Enumerable.Range(0,state.Count).All(j => state.Count - j  != Math.Abs(row - state[j])) )
                {
                    yield return row;
                }
            }
        }

        public IEnumerable<Lists.IList<int>> InitialStates
        {
            get { yield return AlgoritmiaFactory.Lists<int>.Default; }
        }

        public bool IsFinal(Lists.IList<int> state)
        {
            return state.Count == _n;
        }

    }

    public class NQueensEnumerator1
    {
        protected readonly Func<Lists.IList<int>> _createList;
        protected Lists.IList<int> _decisions;
        protected NQueensStateSpace1 _space;

        public NQueensEnumerator1(Func<Lists.IList<int>> createList = null )
        {
            _createList = createList ?? (() => AlgoritmiaFactory.Lists<int>.Default);
        }

        public Lists.IList<int> FirstSolution(NQueensStateSpace1 space)
        {
            return Enumerate(space).First();
        }

        public IEnumerable<Lists.IList<int>> Enumerate(NQueensStateSpace1 space)
        {
            _decisions = _createList();
            _space = space;
            var initial = space.InitialStates.First();
            if (initial != null)
                foreach (var result in Backtracking(initial))
                    yield return result;
        }

        protected virtual IEnumerable<Lists.IList<int>> Backtracking(Lists.IList<int> state)
        {
            if (_space.IsFinal(state))
            {
                yield return state;
            }
            foreach(var decision in _space.Decisions(state))
            {
                _decisions.Add(decision);
                var successor = _space.Decide(state, decision);
                foreach(var result in Backtracking(successor))
                {
                    yield return result;
                }
                _decisions.RemoveAt(_decisions.Count - 1);
            }
        }
    }
}
